#include "warray/warray_dft.h"
#include "test.h"

#include <iostream>
#include <limits>

#include "custom_type.h"
#include "helper.h"

typedef std::complex<float> cfloat;
typedef std::complex<double> cdouble;

typedef boost::mpl::list<float,double> scalar_types;
typedef boost::mpl::list<cfloat,cdouble> complex_types;
typedef boost::mpl::list<float,double,cfloat,cdouble> all_types;

BOOST_AUTO_TEST_CASE_TEMPLATE( naive_dft_1d, T, scalar_types )
{
	typedef std::complex<T> complex;
	warp::detail::naive_nd_dft_implementor<T,1> dft1d1e(warp::extents[1]);
	warp::detail::naive_nd_dft_implementor<T,1> dft1d3e(warp::extents[3]);
	warp::detail::naive_nd_dft_implementor<T,1> dft1d4e(warp::extents[4]);
	warp::detail::naive_nd_dft_implementor<T,1> dft1d5e(warp::extents[5]);
	warp::detail::naive_nd_dft_implementor<T,1> dft1d16e(warp::extents[16]);
	
	// 1 element transforms
	complex inputs_1e[][1] = { {complex(0,0)},{complex(1,0)},{complex(2,1)} };
	const int n_1e = sizeof(inputs_1e)/sizeof(inputs_1e[0]);
	complex output_1e[n_1e][1] = { {complex(0,0)},{complex(0,0)},{complex(0,0)} };
	complex answer_1e[n_1e][1] = { {complex(0,0)},{complex(1,0)},{complex(2,1)} };
	
	for(int i=0;i<n_1e;++i)
	{
		dft1d1e.fwd(
			warp::arraynd<complex,1>(warp::extents[1],inputs_1e[i]),
			warp::arraynd<complex,1>(warp::extents[1],output_1e[i]));
		for(int j=0;j<1;++j)
		{
			BOOST_CHECK_CLOSE(output_1e[i][j].real(),answer_1e[i][j].real(),0.0001);
			BOOST_CHECK_CLOSE(output_1e[i][j].imag(),answer_1e[i][j].imag(),0.0001);
		}
	}
	
	// 3 element transforms
	complex inputs_3e[][3] = { 
		{complex(0,0),complex(0,0),complex(0,0)},
		{complex(1,0),complex(2,0),complex(3,0)},
		{complex(1,2),complex(3,-4),complex(5,6)} 
		};
	const int n_3e = sizeof(inputs_3e)/sizeof(inputs_3e[0]);
	complex output_3e[n_3e][3];
	complex answer_3e[n_3e][3] = { 
		{complex(0,0),complex(0,0),complex(0,0)},
		{complex(6,0),complex(-1.5,0.86602540378444),complex(-1.5,-0.86602540378444)},
		{complex(9,4),complex(-11.66025403784439,2.7320580756888),complex(5.66025403784439,-0.73205080756888)} 
		};
	
	for(int i=0;i<n_3e;++i)
	{
		dft1d3e.fwd(
			warp::arraynd<complex,1>(warp::extents[3],inputs_3e[i]),
			warp::arraynd<complex,1>(warp::extents[3],output_3e[i]));
		//std::cout << "I:" << inputs_3e[i][0] << inputs_3e[i][1] << inputs_3e[i][2] << std::endl;
		//std::cout << "O:" << output_3e[i][0] << output_3e[i][1] << output_3e[i][2] << std::endl;
		//std::cout << "A:" << answer_3e[i][0] << answer_3e[i][1] << answer_3e[i][2] << std::endl;
		for(int j=0;j<3;++j)
		{
			BOOST_CHECK_CLOSE(output_3e[i][j].real(),answer_3e[i][j].real(),0.01);
			BOOST_CHECK_CLOSE(output_3e[i][j].imag(),answer_3e[i][j].imag(),0.01);
		}
	}
	
	// 4 element transforms
	complex inputs_4e[][4] = { 
		{complex(0,0),complex(0,0),complex(0,0),complex(0,0)},
		{complex(1,0),complex(2,0),complex(3,0),complex(4,0)},
		{complex(1,2),complex(3,-4),complex(5,6),complex(7,-8)} 
		};
	const int n_4e = sizeof(inputs_4e)/sizeof(inputs_4e[0]);
	complex output_4e[n_4e][4];
	complex answer_4e[n_4e][4] = { 
		{complex(0,0),complex(0,0),complex(0,0),complex(0,0)},
		{complex(10,0),complex(-2,2),complex(-2,0),complex(-2,-2)},
		{complex(16,-4),complex(0,0),complex(-4,20),complex(-8,-8)} 
		};
	
	for(int i=0;i<n_4e;++i)
	{
		dft1d4e.fwd(
			warp::arraynd<complex,1>(warp::extents[4],inputs_4e[i]),
			warp::arraynd<complex,1>(warp::extents[4],output_4e[i]));
		for(int j=0;j<4;++j)
		{
			// We need these because BOOST_CHECK_CLOSE does not handle zero very well at all
			if(std::abs(output_4e[i][j].real()-answer_4e[i][j].real()) > 0.0001)
				BOOST_CHECK_CLOSE(output_4e[i][j].real(),answer_4e[i][j].real(),0.01);
			if(std::abs(output_4e[i][j].imag()-answer_4e[i][j].imag()) > 0.0001)
				BOOST_CHECK_CLOSE(output_4e[i][j].imag(),answer_4e[i][j].imag(),0.01);
		}
	}
	
	// 5 element transforms
	complex inputs_5e[][5] = { 
		{complex(0,0),complex(0,0),complex(0,0),complex(0,0),complex(0,0)},
		{complex(1,0),complex(2,0),complex(3,0),complex(4,0),complex(5,0)},
		{complex(1,2),complex(3,-4),complex(5,6),complex(7,-8),complex(9,10)} 
		};
	const int n_5e = sizeof(inputs_5e)/sizeof(inputs_5e[0]);
	complex output_5e[n_5e][5];
	complex answer_5e[n_5e][5] = { 
		{complex(0,0),complex(0,0),complex(0,0),complex(0,0),complex(0,0)},
		{complex(15,0),complex(-2.5,3.44095480117793),complex(-2.5,0.81229924058227),complex(-2.5,-0.81229924058227),complex(-2.5,-3.44095480117793)},
		{complex(25,6),complex(-10.08579769603753,12.35404555735545),complex(-26.54378476022677,-1.84753747383505),complex(16.54378476022677,-5.09673443616411),complex(0.08579769603753,-1.40977364735629)} 
		};
	
	for(int i=0;i<n_5e;++i)
	{
		dft1d5e.fwd(
			warp::arraynd<complex,1>(warp::extents[5],inputs_5e[i]),
			warp::arraynd<complex,1>(warp::extents[5],output_5e[i]));
		for(int j=0;j<5;++j)
		{
			BOOST_CHECK_CLOSE(output_5e[i][j].real(),answer_5e[i][j].real(),0.01);
			BOOST_CHECK_CLOSE(output_5e[i][j].imag(),answer_5e[i][j].imag(),0.01);
		}
	}
	
	// 16 element transforms
	complex inputs_16e[][16] = { 
		{complex(0,0),complex(0,0),complex(0,0),complex(0,0),
		 complex(0,0),complex(0,0),complex(0,0),complex(0,0),
		 complex(0,0),complex(0,0),complex(0,0),complex(0,0),
		 complex(0,0),complex(0,0),complex(0,0),complex(0,0)},
		{complex(1,0),complex(2,0),complex(3,0),complex(4,0),
		 complex(5,0),complex(6,0),complex(7,0),complex(8,0),
		 complex(9,0),complex(10,0),complex(11,0),complex(12,0),
		 complex(13,0),complex(14,0),complex(15,0),complex(16,0)},
		{complex(1,2),complex(3,-4),complex(5,6),complex(7,-8),
		 complex(9,10),complex(11,-12),complex(13,14),complex(15,-16),
		 complex(17,18),complex(19,-20),complex(21,22),complex(23,-24),
		 complex(25,26),complex(27,-28),complex(29,30),complex(31,-32)} 
		};
	const int n_16e = sizeof(inputs_16e)/sizeof(inputs_16e[0]);
	complex output_16e[n_16e][16];
	complex answer_16e[n_16e][16] = { 
		{complex(0,0),complex(0,0),complex(0,0),complex(0,0),
		 complex(0,0),complex(0,0),complex(0,0),complex(0,0),
		 complex(0,0),complex(0,0),complex(0,0),complex(0,0),
		 complex(0,0),complex(0,0),complex(0,0),complex(0,0)},
		{complex(136,0),complex(-8,40.21871593700678),complex(-8,19.31370849898476),complex(-8,11.97284610132391),
		 complex(-8,8),complex(-8,5.34542910335439),complex(-8,3.31370849898476),complex(-8,1.59129893903726),
		 complex(-8,0),complex(-8,-1.59129893903726),complex(-8,-3.31370849898476),complex(-8,-5.34542910335439),
		 complex(-8,-8),complex(-8,-11.97284610132391),complex(-8,-19.31370849898476),complex(-8,-40.21871593700678)},
		{complex(256.000000000000000,-16.000000000000000),
		 complex(-12.817402121925474,+64.437431874013569),
		 complex(-9.372583002030478,+22.627416997969522),
		 complex(-5.309141793291221,+7.945692202647825),
		 complex(0.000000000000000,+0.000000000000000),
		 complex(7.945692202647825,-5.309141793291220),
		 complex(22.627416997969522,-9.372583002030478),
		 complex(64.437431874013569,-12.817402121925474),
		 complex(-16.000000000000000,+272.000000000000000),
		 complex(-96.437431874013569,-19.182597878074525),
		 complex(-54.627416997969519,-22.627416997969522),
		 complex(-39.945692202647820,-26.690858206708782),
		 complex(-32.000000000000000,-32.000000000000000),
		 complex(-26.690858206708782,-39.945692202647827),
		 complex(-22.627416997969522,-54.627416997969519),
		 complex(-19.182597878074525,-96.437431874013569)} 
		};
	
	for(int i=0;i<n_16e;++i)
	{
		dft1d16e.fwd(
			warp::arraynd<complex,1>(warp::extents[16],inputs_16e[i]),
			warp::arraynd<complex,1>(warp::extents[16],output_16e[i]));
		for(int j=0;j<16;++j)
		{
			// We need these because BOOST_CHECK_CLOSE does not handle zero very well at all
			if(std::abs(output_16e[i][j].real()-answer_16e[i][j].real()) > 0.0001)
				BOOST_CHECK_CLOSE(output_16e[i][j].real(),answer_16e[i][j].real(),0.01);
			if(std::abs(output_16e[i][j].imag()-answer_16e[i][j].imag()) > 0.0001)
				BOOST_CHECK_CLOSE(output_16e[i][j].imag(),answer_16e[i][j].imag(),0.01);
		}
	}
}

BOOST_AUTO_TEST_CASE_TEMPLATE( naive_dft_2d, T, scalar_types )
{
	typedef std::complex<T> complex;
	warp::detail::naive_nd_dft_implementor<T,2> dft2d2x4(warp::extents[2][4]);
	
	// 2x4 element transforms
	complex inputs_2x4[][2*4] = { 
		{complex(0,0),complex(0,0),complex(0,0),complex(0,0),
		 complex(0,0),complex(0,0),complex(0,0),complex(0,0)},
		{complex(1,0),complex(2,0),complex(3,0),complex(4,0),
		 complex(5,0),complex(6,0),complex(7,0),complex(8,0)},
		{complex(1,2),complex(3,-4),complex(5,6),complex(7,-8),
		 complex(9,10),complex(11,-12),complex(13,14),complex(15,-16)}  
		};
	const int n_2x4 = sizeof(inputs_2x4)/sizeof(inputs_2x4[0]);
	complex output_2x4[n_2x4][2*4];
	complex answer_2x4[n_2x4][2*4] = { 
		{complex(0,0),complex(0,0),complex(0,0),complex(0,0),
		 complex(0,0),complex(0,0),complex(0,0),complex(0,0)},
		{complex( 36,0),complex(-4,4),complex(-4,0),complex(-4,-4),
		 complex(-16,0),complex( 0,0),complex( 0,0),complex( 0,0)},
		{complex( 64,-8),complex(0,0),complex(-8, 72),complex(-16,-16),
		 complex(-32, 0),complex(0,0),complex( 0,-32),complex(  0,  0)},
		};
	
	for(int i=0;i<n_2x4;++i)
	{
		dft2d2x4.fwd(
			warp::arraynd<complex,2>(warp::extents[2][4],inputs_2x4[i]),
			warp::arraynd<complex,2>(warp::extents[2][4],output_2x4[i]));
		for(int j=0;j<1;++j)
		{
			BOOST_CHECK_CLOSE(output_2x4[i][j].real(),answer_2x4[i][j].real(),0.0001);
			BOOST_CHECK_CLOSE(output_2x4[i][j].imag(),answer_2x4[i][j].imag(),0.0001);
		}
	}
}

BOOST_AUTO_TEST_CASE_TEMPLATE( naive_dft_3d, T, scalar_types )
{
	typedef std::complex<T> complex;
	warp::detail::naive_nd_dft_implementor<T,3> dft3d2x2x4(warp::extents[2][2][4]);
	
	// 2x4 element transforms
	complex inputs_2x2x4[][2*2*4] = { 
		{complex(0,0),complex(0,0),complex(0,0),complex(0,0),
		 complex(0,0),complex(0,0),complex(0,0),complex(0,0),
		 complex(0,0),complex(0,0),complex(0,0),complex(0,0),
		 complex(0,0),complex(0,0),complex(0,0),complex(0,0)},
		{complex(1,0),complex(2,0),complex(3,0),complex(4,0),
		 complex(5,0),complex(6,0),complex(7,0),complex(8,0),
		 complex(9,0),complex(10,0),complex(11,0),complex(12,0),
		 complex(13,0),complex(14,0),complex(15,0),complex(16,0)},
		{complex(1,2),complex(3,-4),complex(5,6),complex(7,-8),
		 complex(9,10),complex(11,-12),complex(13,14),complex(15,-16),
		 complex(17,18),complex(19,-20),complex(21,22),complex(23,-24),
		 complex(25,26),complex(27,-28),complex(29,30),complex(31,-32)} 
		};
	const int n_2x2x4 = sizeof(inputs_2x2x4)/sizeof(inputs_2x2x4[0]);
	complex output_2x2x4[n_2x2x4][2*2*4];
	complex answer_2x2x4[n_2x2x4][2*2*4] = { 
		{complex(0,0),complex(0,0),complex(0,0),complex(0,0),
		 complex(0,0),complex(0,0),complex(0,0),complex(0,0),
		 complex(0,0),complex(0,0),complex(0,0),complex(0,0),
		 complex(0,0),complex(0,0),complex(0,0),complex(0,0)},
		{complex(136,0),complex(-8,8),complex(-8,0),complex(-8,-8),
		 complex(-32,0),complex( 0,0),complex( 0,0),complex( 0, 0),
		 complex(-64,0),complex( 0,0),complex( 0,0),complex( 0, 0),
		 complex(  0,0),complex( 0,0),complex( 0,0),complex( 0, 0)},
		{complex(256,-16),complex(0,0),complex(-16, 272),complex(-32,-32),
		 complex(-64,  0),complex(0,0),complex(  0, -64),complex(  0,  0),
		 complex(-128, 0),complex(0,0),complex(  0,-128),complex(  0,  0),
		 complex(   0, 0),complex(0,0),complex(  0,   0),complex(  0,  0)},
		};
	
	for(int i=0;i<n_2x2x4;++i)
	{
		dft3d2x2x4.fwd(
			warp::arraynd<complex,3>(warp::extents[2][2][4],inputs_2x2x4[i]),
			warp::arraynd<complex,3>(warp::extents[2][2][4],output_2x2x4[i]));
		for(int j=0;j<1;++j)
		{
			BOOST_CHECK_CLOSE(output_2x2x4[i][j].real(),answer_2x2x4[i][j].real(),0.0001);
			BOOST_CHECK_CLOSE(output_2x2x4[i][j].imag(),answer_2x2x4[i][j].imag(),0.0001);
		}
	}
}


